import paddle

#Encoder

class Encoder(paddle.nn.Layer):
    
    def __init__(self,in_channels=3):
        super(Encoder,self).__init__()
        self.conv1 = paddle.nn.Conv2D(3,64,kernel_size=3,stride=3)
        self.conv2 = paddle.nn.Conv2D(64,64,kernel_size=3,stride=3)
        self.conv3 = paddle.nn.Conv2D(64,32,kernel_size=3,stride=3)
        self.conv4 = paddle.nn.Conv2D(32,16,kernel_size=3,stride=3)
        self.relu = paddle.nn.ReLU()
        self.flatten = paddle.nn.Flatten()
        self.linear = paddle.nn.Linear(256,2)
        
    def forward(self,x):
        y = self.conv1(x)
        y = self.relu(y)
        y = self.conv2(y)
        y = self.relu(y)
        y = self.conv3(y)
        y = self.relu(y)
        y = self.conv4(y)
        y = self.relu(y)
        y = self.flatten(y)
        y = self.linear(y)
        return y
    
#Decoder

class Decoder(paddle.nn.Layer):
    def __init__(self,in_channels=3):
        super(Decoder,self).__init__()
        self.linear = paddle.nn.Linear(2,256)
        self.relu = paddle.nn.ReLU()
        self.conv1 =  paddle.nn.Conv2D(16,576,kernel_size=3,stride=1,padding=1)
        self.conv2 =  paddle.nn.Conv2D(64,288,kernel_size=3,stride=1,padding=1)
        self.conv3 =  paddle.nn.Conv2D(32,144,kernel_size=3,stride=1,padding=1)
        self.conv4 =  paddle.nn.Conv2D(16,144,kernel_size=3,stride=1,padding=1)
        self.conv5 =  paddle.nn.Conv2D(16,3,kernel_size=3,stride=1,padding=1)
        
    def forward(self,x):
        y = self.linear(x)
        y = paddle.reshape(y,(x.shape[0],16, 4, 4))
        y = self.conv1(y)
        y = self.relu(y)
        y = paddle.nn.functional.pixel_shuffle(y,3)
        y = self.conv2(y)
        y = self.relu(y)
        y = paddle.nn.functional.pixel_shuffle(y,3)
        y = self.conv3(y)
        y = self.relu(y)
        y = paddle.nn.functional.pixel_shuffle(y,3)
        y = self.conv4(y)
        y = self.relu(y)
        y = paddle.nn.functional.pixel_shuffle(y,3)
        y = self.conv5(y)
        
        return y
    
#VAE
class VAE(paddle.nn.Layer):
    def __init__(self):
        super(VAE,self).__init__()
        self.encoder = Encoder()
        self.decoder = Decoder()
    def forward(self,x):
        y = self.encoder(x)
        y = self.decoder(y)
        return y
    
    
if __name__ == "__main__":
    vae = VAE()
    encoder = Encoder()
    decoder = Decoder()
    
    paddle.summary(vae,(1,3,324,324))
    paddle.summary(encoder,(1,3,324,324))
    paddle.summary(decoder,(1,2))
    
























